#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import time
import argparse
from src.commands import asset_command
from src.core.database import DatabaseConnector
from src.commands.asset_command import AssetCommand
from src.commands.owner_command import OwnerCommand
from src.commands.ip_command import IPCommand
from src.commands.domain_command import DomainCommand
from src.commands.ipdomain_command import IPDomainCommand
from src.core.help_generator import HelpGenerator
from src.core.dbhelper import DBHelper

class AssetCollectionSystem:
    """资产收集系统主类"""
    
    def __init__(self):
        """初始化资产收集系统"""
        self.db_connector = DatabaseConnector()
        self.db_helper = DBHelper(self.db_connector)
        
        self.parser = argparse.ArgumentParser(description='资产数据库管理工具')
        subparsers = self.parser.add_subparsers(dest='command', help='可用命令')
        # help命令
        help_parser = subparsers.add_parser('help', help='显示帮助信息')
        # add_target命令
        add_target_parser = subparsers.add_parser('add_target', help='添加目标(IP或域名)')
        add_target_parser.add_argument('target', help='目标IP或域名，多个用逗号分隔')
        add_target_parser.add_argument('--owner', help='目标所属单位ID或名称')
        # import_target命令
        # import_target_parser = subparsers.add_parser('import_target', help='导入目标')
        # import_target_parser.add_argument('target_file', help='目标文件路径')
        # import_target_parser.add_argument('--owner-column', help='单位列名，内容可以是单位名称或者系统中的单位ID，默认为owner，或者第一列', default='owner')
        # import_target_parser.add_argument('--target-column', help='目标列名，内容可以是IP或者域名，默认为target，或者第二列', default='target')
        # asset命令
        self.asset_command = AssetCommand(self.db_connector, subparsers)
        self.owner_command = OwnerCommand(self.db_connector, subparsers)
        self.ip_command = IPCommand(self.db_connector, subparsers)
        self.domain_command = DomainCommand(self.db_connector, subparsers)
        self.ipdomain_command = IPDomainCommand(self.db_connector, subparsers)
 
    
    def handle_command_line_args(self, args):
        """处理命令行参数"""
        if args.command == 'help':
            self.parser.print_help()
        elif args.command == 'asset':
            self.asset_command.handle_command(args)
        elif args.command == 'owner':
            self.owner_command.handle_command(args)
        elif args.command == 'domain':
            self.domain_command.handle_command(args)
        elif args.command == 'ipdomain':
            self.ipdomain_command.handle_command(args)
        elif args.command == 'ip':
            self.ip_command.handle_command(args)
        elif args.command == 'add_target':
            self.add_target(args.target, args.owner)
        elif args.command == 'import_target':
            self.import_target(args.target_file, args.owner_column, args.target_column)
        else:
            print(f"未知命令: {args.command}")
            self.parser.print_help()
    
    def add_target(self, target, owner_identifier):
        """添加目标"""
        targets = target.split(',')
        self.db_helper.add_ip_or_domain(targets, owner_identifier)

    def import_target(self, target_file, owner_column, target_column):
        """导入目标"""
        pass

    def initialize(self):
        """初始化系统"""
        try:
            # 连接数据库
            self.db_connector.connect()
            # 检查数据库结构
            self.db_connector.check_database_structure()
            print("数据库连接成功！")
            
            # 获取asset_parser
            asset_parser = None
            for action in self.parser._actions:
                if isinstance(action, argparse._SubParsersAction) and 'asset' in action.choices:
                    asset_parser = action.choices['asset']
                    break
            
            # 初始化命令处理器
            self.asset_command = AssetCommand(self.db_connector, asset_parser)
            self.owner_command = OwnerCommand(self.db_connector)
            self.ip_command = IPCommand(self.db_connector)
            self.domain_command = DomainCommand(self.db_connector)
            self.ipdomain_command = IPDomainCommand(self.db_connector)
        except Exception as e:
            print(f"系统初始化失败: {e}")
            return False
        
        return True


def main():
    """主函数"""
    # 确保中文正常显示
    import io
    import sys
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    
    # 创建系统实例
    system = AssetCollectionSystem()
    
    # 初始化系统
    # if not system.initialize():
    #     print("系统初始化失败，程序退出")
    #     sys.exit(1)
    
    # 解析命令行参数
    args = system.parser.parse_args()
    
    # 处理命令行参数
    system.handle_command_line_args(args)
    
    # 关闭数据库连接
    try:
        system.db_connector.disconnect()
    except Exception as e:
        print(f"关闭数据库连接出错: {e}")
    
    sys.exit(0)

if __name__ == "__main__":
    main()